package com.yilin.gateway.filter_factory;

import com.google.gson.reflect.TypeToken;
import com.yilin.gateway.core.EncryptUtils;
import com.yilin.gateway.core.JsonUtil;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;

public class EncryptResponseBodyGatewayFilterFactory extends AbstractGatewayFilterFactory {
    private final static String TOKEN = "token";

    public EncryptResponseBodyGatewayFilterFactory() {
    }

    @Override
    public GatewayFilter apply(Object config) {
        return new EncryptResponseGatewayFilter();
    }


    public static class EncryptResponseGatewayFilter implements GatewayFilter, Ordered {
        @Override
        public int getOrder() {
            return -2;
        }

        @Override
        public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
            return chain.filter(exchange.mutate().response(decorate(exchange)).build());
        }

        @SuppressWarnings("unchecked")
        private ServerHttpResponse decorate(ServerWebExchange exchange) {
            return new ServerHttpResponseDecorator(exchange.getResponse()) {
                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                    String originalResponseContentType = exchange.getAttribute("ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR");
                    HttpHeaders httpHeaders = new HttpHeaders();
                    httpHeaders.add(HttpHeaders.CONTENT_TYPE, originalResponseContentType);
                    String tokenID = exchange.getRequest().getHeaders().getFirst("token");
                    //判断返回的数据是否是JSON
                    String contentType = exchange.getResponse().getHeaders().getFirst("Content-Type");
                    if (contentType == null || !contentType.contains("application/json")) {
                        return getDelegate().writeWith(body);
                    }
                    ClientResponse clientResponse = ClientResponse.create(Objects.requireNonNull(exchange.getResponse().getStatusCode()))
                            .headers(headers -> headers.putAll(httpHeaders)).body(Flux.from(body)).build();
                    //修改body
                    Mono<String> modifiedBody = clientResponse.bodyToMono(String.class).flatMap(originalBody -> modifyBody(tokenID).apply(exchange, Mono.just(originalBody)));
                    BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
                    CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, exchange.getResponse().getHeaders());
                    return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
                        Flux<DataBuffer> messageBody = outputMessage.getBody();
                        HttpHeaders headers = getDelegate().getHeaders();
                        if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
                            messageBody = messageBody.doOnNext(data -> headers.setContentLength(data.readableByteCount()));
                        }
                        return getDelegate().writeWith(messageBody);
                    }));
                }

                /**
                 * 修改body
                 * @return apply 返回Mono<String>，数据是修改后的body
                 */
                private BiFunction<ServerWebExchange, Mono<String>, Mono<String>> modifyBody(String token) {
                    return (exchange, json) -> {
                        AtomicReference<String> result = new AtomicReference<>();
                        json.subscribe(value -> {
                            Map<String, Object> map = JsonUtil.jsonToObject(value, new TypeToken<Map<String, Object>>() {});
                            String headToken = token;
                            if (map != null && map.get(TOKEN) != null) headToken = (String) map.get(TOKEN);
                            if (map != null && map.containsKey("data")) {
                                if (map.get("data") != null) {
                                    String jsonData = JsonUtil.objectToJson(map.get("data"));
                                    map.put("encData", EncryptUtils.aesEncrypt(jsonData, headToken));
//                                    map.remove("data");
                                    //map.put("data",EncryptUtils.aesEncrypt(jsonData,headToken));
                                }
                            }
                            result.set(JsonUtil.objectToJson(map));
                        }, Throwable::printStackTrace);
                        return Mono.just(result.get());
                    };
                }

                @Override
                public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
                    return writeWith(Flux.from(body).flatMapSequential(p -> p));
                }
            };
        }
    }
}
